import sys
from .ray_manager import RayManager
from calc_manager.file import FileOperatorHelper
from calc_manager.algorithm.graph import AlgoGraph
from calc_manager.repo.sh_res_repo import (ShDcAlgoRepository,
                                           ShDcOperatorRepository)
from calc_manager.algorithm.param import ParameterMappingService
from calc_manager.algorithm.algo import (Algorithm, AlgorithmInfo, OperatorComponentToOperatorInfoAdapter)
from calc_manager.executor.exec import build_algo_tasks
from calc_manager.ray.remote_task import RayAlgoOperatorTaskFactory

import pickle
import base64
import os

def get_ray_task_factory_to_pickle_b64(factory_task: RayAlgoOperatorTaskFactory) -> str:
    task_factory_serializeable = pickle.dumps(factory_task)
    enc = base64.b64encode(task_factory_serializeable)
    return enc.decode('utf-8')

def get_ray_task_factory_to_pickle_b64(factory_task: RayAlgoOperatorTaskFactory) -> str:
    task_factory_serializeable = pickle.dumps(factory_task)
    enc = base64.b64encode(task_factory_serializeable)
    return enc.decode('utf-8')

def get_ray_task_factory_from_pickle_b64(pickle_b64: str) -> RayAlgoOperatorTaskFactory:
    pickle_b64 = base64.b64decode(pickle_b64)
    return pickle.loads(pickle_b64)


def main():
    try:
        encoded_data = os.getenv("ENCODED_DATA", "")
        remote_task_algo = get_ray_task_factory_from_pickle_b64(encoded_data)
        ray_tasks = build_algo_tasks(remote_task_algo, remote_task_algo.algo)
        manager = RayManager("", {},)
        res_ray, ctx = manager.invoke_tasks(ray_tasks)
        print(res_ray)
    except Exception as err:
        print(err)
        sys.exit(1)
